import os.path

import pandas as pd
from pypinyin import lazy_pinyin  # 新增导入
import config_manager



class BastionLoader:
    def __init__(self, columns=None):
        self.excel_path = config_manager.ConfigManager().get_login_path()
        if self.excel_path is None or self.excel_path == "":
            config_manager.ConfigManager().set_login_path("login.xlsx")
            self.excel_path = "login.xlsx"
        self.data = None
        self.columns = columns # 保存列名，如果提供了的话
    
    def load_data(self):
        """从Excel文件加载堡垒机数据并排序"""
        # 使用传入的columns，如果没有则使用默认硬编码列
        cols_to_use = self.columns if self.columns is not None else ["路段名称", "堡垒机IP", "用户名", "登录密码", "堡垒机品牌", "备注"]
        try:
            self.data = pd.read_excel(self.excel_path)
            self.data = self.data.dropna(subset=["路段名称"])
            # 确保只选择存在的列，避免KeyError
            existing_cols = [col for col in cols_to_use if col in self.data.columns]
            self.data = self.data[existing_cols]
            self.data = self.data.replace({"" : pd.NA })
            self.data = self.data.dropna(subset= ["路段名称"])
            self.data.fillna("", inplace=True)
            # 修改为按中文拼音排序
            self.data = self.data.iloc[pd.Series(lazy_pinyin(str(x)) for x in self.data['路段名称']).argsort()]
            # print(self.data)
            return True
        except Exception as e:
            print(f"加载Excel文件失败: {e}")
            # 返回只包含表头的空DataFrame
            self.data = pd.DataFrame(columns=cols_to_use)
            return False
    
    def get_sorted_data(self):
        """获取排序后的堡垒机数据"""
        if self.data is not None:
            return self.data
        return None
    
    def search_data(self, field, keyword):
        """根据字段和关键词搜索堡垒机数据"""
        if self.data is None:
            return None
        return self.data[self.data[field].str.contains(keyword, case=False, na=False)]


class ServersLoader:
    def __init__(self, columns=None):
        login_path = config_manager.ConfigManager().get_servers_path()
        if  login_path is None:
            config_manager.ConfigManager().set_servers_path("servers_info.xlsx")
            login_path = "servers_info.xlsx"

        self.excel_path = login_path
        self.data = None
        self.columns = columns # 保存列名，如果提供了的话

    def load_data(self):
        """从Excel文件加载堡垒机数据并排序"""
        # 使用传入的columns，如果没有则使用默认硬编码列
        cols_to_use = self.columns if self.columns is not None else ["分类", "路段", "堡垒机", "跳机", "密码", "前置机", "密码", "超融合LinuxIP",
                "Origin_IP", "Gate_IP", "超融合Linux密码", "超融合Windows密码", "站级密码"]
        try:
            self.data = pd.read_excel(self.excel_path)
            self.data = self.data.dropna(subset=["路段"])
            # 确保只选择存在的列，避免KeyError
            existing_cols = [col for col in cols_to_use if col in self.data.columns]
            self.data = self.data[existing_cols]
            # 新增：去除所有字符串中的换行符（\n）
            self.data = self.data.applymap(lambda x: x.replace('\n', '') if isinstance(x, str) else x)
            self.data.fillna("", inplace=True)
            # 修改为按中文拼音排序
            self.data = self.data.iloc[pd.Series(lazy_pinyin(str(x)) for x in self.data['路段']).argsort()]
            # print(self.data)
            return True
        except Exception as e:
            print(f"加载Excel文件失败: {e}")
            # 返回只包含表头的空DataFrame
            self.data = pd.DataFrame(columns=cols_to_use)
            return False

    def get_sorted_data(self):
        """获取排序后的堡垒机数据"""
        if self.data is not None:
            return self.data
        return None

    def search_data(self, field, keyword):
        """根据字段和关键词搜索堡垒机数据"""
        if self.data is None:
            return None
        return self.data[self.data[field].str.contains(keyword, case=False, na=False)]

if __name__ == "__main__":
    loader = BastionLoader()
    if loader.load_data():
        print(loader.get_sorted_data())